package com.ljt.aicodehelp.ai;

import com.ljt.aicodehelp.ai.guardrail.SensitiveWordGuardRail;
import com.ljt.aicodehelp.ai.tool.InterviewQuestionTool;
import dev.langchain4j.community.model.dashscope.QwenStreamingChatModel;
import dev.langchain4j.mcp.McpToolProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.service.AiServices;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.List;

@Configuration
public class AiCodeHelperFactory {

    @Bean
    public AiCodeHelperService aiCodeHelperService(
            ChatModel myQwenChatModel,
            ContentRetriever contentRetriever,
            McpToolProvider mcpToolProvider,
            QwenStreamingChatModel qwenStreamingChatModel
    ){
        // 会话记忆，只记忆10条上下文
        var messageWindowChatMemory = MessageWindowChatMemory.builder()
                .maxMessages(10)
                .build();

        // chatMemory和chatMemoryProvider只能二选一
        return AiServices
                .builder(AiCodeHelperService.class)
                // .chatMemory(messageWindowChatMemory)
                .chatModel(myQwenChatModel)
                .streamingChatModel(qwenStreamingChatModel)
                .chatMemoryProvider(memoryId ->
                        MessageWindowChatMemory.withMaxMessages(10))
                .contentRetriever(contentRetriever)
                .tools(List.of(new InterviewQuestionTool()))
                .toolProvider(mcpToolProvider)
                .inputGuardrails(new SensitiveWordGuardRail())
                .build();
    }
}
